import os
import uuid
from chromadb import Client, Settings
import dashscope
from dashscope import TextEmbedding

# 配置信息
DASHSCOPE_API_KEY = 'sk-4a106bdba4ce4b01b836651877585128'
EMBEDDING_MODEL = 'text-embedding-v4'
COLLECTION_NAME = 'rag_docs'

# 设置API Key
dashscope.api_key = DASHSCOPE_API_KEY

# 创建纯内存模式的Chroma客户端
print('创建纯内存模式的Chroma客户端...')
client = Client(Settings(
    chroma_db_impl='duckdb+memory',
    anonymized_telemetry=False
))

# 创建或获取集合
collection = client.get_or_create_collection(name=COLLECTION_NAME)
print(f'集合中文档数量: {collection.count()}')

# 定义向量化函数
def get_embedding(text):
    try:
        truncated_text = text.strip()[:5000] if text else '空内容'
        response = TextEmbedding.call(
            model=EMBEDDING_MODEL,
            input=truncated_text
        )
        if response.status_code == 200:
            return response.output['embeddings'][0]['embedding']
        else:
            print(f'嵌入失败: {response.status_code} - {response.message}')
            return [0.0] * 1024  # 返回零向量
    except Exception as e:
        print(f'调用嵌入模型失败: {e}')
        return [0.0] * 1024  # 返回零向量

# 检查文件是否存在并读取内容
file_path = r'c:\Users\25188\Downloads\git\backend\uploads\txt_backup.txt'
print(f'\n检查文件: {file_path}')
if os.path.exists(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        file_content = f.read()
    print(f'文件内容: {file_content}')

    # 手动添加文件内容到向量数据库
    print('\n手动添加文件内容到向量数据库...')
    chunks = [{'content': file_content, 'metadata': {'source': 'txt_backup.txt'}}]
    texts = [chunk['content'] for chunk in chunks]
    metadatas = [chunk['metadata'] for chunk in chunks]
    ids = [str(uuid.uuid4()) for _ in range(len(texts))]

    # 获取嵌入向量
    print('获取嵌入向量...')
    embeddings = [get_embedding(text) for text in texts]

    # 添加到集合
    collection.add(
        embeddings=embeddings,
        documents=texts,
        metadatas=metadatas,
        ids=ids
    )
    print(f'成功添加 {len(texts)} 个文本块到向量数据库')
    print(f'更新后集合中文档数量: {collection.count()}')
else:
    print('文件不存在')

# 测试检索功能
query = '你是谁？'
print(f'\n测试检索: "{query}"')
query_embedding = get_embedding(query)

results = collection.query(
    query_embeddings=[query_embedding],
    n_results=3
)

# 处理检索结果
print(f'检索到 {len(results["documents"][0])} 个结果:')
for i, (doc, meta, dist) in enumerate(zip(results['documents'][0], results['metadata'][0], results['distances'][0]), 1):
    print(f'结果 {i}:')
    print(f'  内容: {doc}')
    print(f'  元数据: {meta}')
    print(f'  相似度距离: {float(dist):.4f}')

print('\n检查完成。如果检索到结果，说明RAG可以引用上传的文件内容。')